"""
边侧节点通信相关的数据类型定义
"""

from typing import Dict, List, Optional, Any
from datetime import datetime
from pydantic import BaseModel, Field
from enum import Enum


class EdgeNodeStatus(str, Enum):
    """边侧节点状态枚举"""
    HEALTHY = "healthy"
    UNHEALTHY = "unhealthy" 
    ERROR = "error"
    OFFLINE = "offline"


class EdgeHealthStatus(BaseModel):
    """边侧节点健康状态"""
    node_id: str
    status: EdgeNodeStatus
    details: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    timestamp: str


class EdgeSampleInfo(BaseModel):
    """边侧节点样本信息"""
    node_id: str
    task_name: str
    collected_sample_count: int
    uncollected_sample_count: int
    uninferred_sample_count: int
    total_sample_count: int
    difficulty_threshold_stats: Dict[str, int]
    timestamp: str


class EdgeAggregationRequest(BaseModel):
    """边侧样本聚合请求"""
    difficulty_threshold: Optional[float] = Field(None, ge=0.0, le=1.0)
    time_start: Optional[str] = None
    time_end: Optional[str] = None
    max_samples: Optional[int] = Field(1000, gt=0, le=10000)


class EdgeAggregationResponse(BaseModel):
    """边侧样本聚合响应"""
    success: bool
    node_id: str
    message: str
    selected_samples: List[str]
    sample_count: int
    upload_info: Optional[Dict[str, Any]] = None
    request_params: Dict[str, Any]
    timestamp: str


class EdgeUploadStatus(BaseModel):
    """边侧上传状态"""
    node_id: str
    task_id: str
    status: str
    details: Dict[str, Any]
    timestamp: str


class EdgeModelDeployment(BaseModel):
    """边侧模型部署配置"""
    model_id: str
    model_url: str
    model_path: Optional[str] = None
    force_update: bool = False
    backup_current: bool = True


class EdgeModelInfo(BaseModel):
    """边侧模型信息"""
    node_id: str
    model_loaded: bool
    current_model_id: Optional[str] = None
    model_path: Optional[str] = None
    model_metadata: Optional[Dict[str, Any]] = None
    timestamp: str


class EdgeInferenceResult(BaseModel):
    """边侧推理结果"""
    node_id: str
    success: bool
    processed_samples: int
    inference_time: Optional[float] = None
    error: Optional[str] = None
    timestamp: str


class EdgeSystemStatus(BaseModel):
    """边侧系统状态"""
    node_id: str
    cpu_usage: float
    memory_usage: float
    disk_usage: float
    network_status: str
    service_status: Dict[str, str]
    timestamp: str


class EdgeOperationResult(BaseModel):
    """边侧操作结果的通用格式"""
    success: bool
    node_id: str
    operation: str
    data: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    timestamp: str


class EdgeBatchOperationResult(BaseModel):
    """边侧批量操作结果"""
    operation: str
    total_nodes: int
    successful_nodes: List[str]
    failed_nodes: List[str] 
    missing_nodes: List[str]
    results: Dict[str, Any]  # 修改为Any类型以支持不同类型的结果
    timestamp: str 